package geppetto.cat.models;


import geppetto.cat.corpus.BilingualCorpus;

import java.io.PrintStream;



/**
 * Implements IBM Model 2
 * @author javg
 *
 */
public class M2 extends M1 {

	DiagonalDistanceDistortionTable _countDist;
	DiagonalDistanceDistortionTable _distortion;
	int distortionBuckets = 30;

	M2() {

	}

	public M2(BilingualCorpus corpus) {
		super(corpus);
		_distortion = new DiagonalDistanceDistortionTable(distortionBuckets);
		_distortion.initialize();
	}

	public M2(BilingualCorpus corpus, SparseTranslationTable tb) {
		super(corpus, tb);
		_distortion = new DiagonalDistanceDistortionTable(distortionBuckets);
		_distortion.initialize();
	}

	public void printStamp(PrintStream file){
		super.printStamp(file);
		file.println("Distortion Nr Buckets: " + distortionBuckets);
	}
	
	// ///// Save and load models from file
	public void saveModel(String directory) {
		super.saveModel(directory);
		_distortion.saveTable(directory);
	}

	public static M2 loadModel(BilingualCorpus corpus, String directory) {
		M2 model = new M2();
		model._corpus = corpus;
		if (!corpus.checkDescription(directory)) {
			System.out.println("Corpus is not the same");
			System.exit(1);
		}
		model._tb = SparseTranslationTable.LoadTranslationTable(corpus,
				directory);
		model._distortion = new DiagonalDistanceDistortionTable(directory);
		// _tb.printSmallTable(System.out);
		model._trained = true;
		return model;
	}

	public String getName() {
		return "IBM M2";
	}

	public DiagonalDistanceDistortionTable getDistortioTable() {
		return _distortion;
	}

	public double calculateProbability(int sourceWord, int foreignWord,
			int sourcePos, int foreignPos, int sourceLen, int foreignLen) {
		return _tb.getProbability(sourceWord, foreignWord)
				* _distortion.getProbability(sourcePos, foreignPos, sourceLen,
						foreignLen) * (1 - 1.0 / (1 + sourceLen));
	}

	public double getNullProbability(int foreignWord, int foreignPos,
			int sourceLen, int foreignLen) {
		return _tb.getNullProbability(foreignWord) / (1 + sourceLen);
	}

	public void clearCounts() {
		_count.initializeToSmoothingValue();
		_countDist.clear();
	}

	public void addCounts(int[] s, int[] f, double[][] posteriors) {
		for (int foreign = 0; foreign < f.length; foreign++) {
			for (int source = 0; source < s.length; source++) {
				_count.addToCount(s[source], f[foreign],
						posteriors[source][foreign]);
				_countDist.setProbability(s[source], f[foreign], s.length,
						f.length, _countDist.getProbability(s[source],
								f[foreign], s.length, f.length)
								+ (float) posteriors[source][foreign]);
			}
			_count.addToNullCount(f[foreign], posteriors[s.length][foreign]);
			// _countDist.(f[foreign], (float)posteriors[s.length][foreign]);
		}
	}

	public void finalizeTrain() {
		_count = null;
		_countDist = null;
	}

	public void initializeTrain() {
		_count = new SparseCountTable(_corpus);
		_countDist = new DiagonalDistanceDistortionTable(distortionBuckets);
		_countDist.initialize();
	}

	public void mStep() {
		_distortion.clear();
		for (int i = 0; i < _distortion.getNRBuckets(); i++) {
			_distortion.setBucketProbability(i, _countDist
					.getBucketProbability(i));
		}
		updateTranslationProbabilities(_tb, _count);
	}

}
